import argparse
from tensorflow.keras.models import load_model
import sys
import os

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..')))

# 添加项目根目录到 sys.path（假设 predict.py 在 PyProject/ProfileSelection/deepcrossing 下）
# from PyProject.ProfileSelection.dataprocess.criteo import create_criteo_tests
# from PyProject.ProfileSelection.deepcrossing.modules import ResidualUnits
# from model import DeepCrossing
from ProfileSelection.dataprocess.criteo import create_criteo_tests
from ProfileSelection.deepcrossing.modules import ResidualUnits
from ProfileSelection.deepcrossing.model import DeepCrossing


def predict_with_model(test_file, model_path='saved_model/deep_crossing_model', embed_dim=8, sample_num=10, batch_size=10, profile_num=10):
    """
    通过加载训练好的模型，并用指定的 test.txt 文件进行预测。

    :param test_file: str, 测试数据的文件路径
    :param model_path: str, 模型保存路径（默认为 'saved_model/deep_crossing_model'）
    :param embed_dim: int, 嵌入维度（默认为 8）
    :param sample_num: int, 需要读取的样本数量
    :param batch_size: int, 批量大小（默认为 10）
    :param profile_num: int, 简历批量大小（默认为 10）
    :return: None, 打印预测结果
    """
    # ========== 读取测试数据 ==========
    feature_columns, test = create_criteo_tests(file=test_file,
                                               embed_dim=embed_dim,
                                               sample_num=sample_num)
    test_X, test_y = test  # 如果 test_y 不可用，也可以只使用 test_X 进行预测

    # ========== 加载已训练模型 ==========
    model = load_model(model_path,
                       custom_objects={'DeepCrossing': DeepCrossing, 'ResidualUnits': ResidualUnits})

    print("模型已加载，开始预测...")

    # ========== 执行预测 ==========
    predictions = model.predict(test_X, batch_size=batch_size)

    # ========== 输出前10个预测结果 ==========
    print("Top predict results:")
    print(predictions[:profile_num])


def main():


    """
    主函数：用于处理命令行参数并调用 predict_with_model 方法。
    """
    # 设置命令行参数解析
    parser = argparse.ArgumentParser(description="使用 DeepCrossing 模型进行预测")
    parser.add_argument('test_file', type=str, help="测试数据文件的路径")
    parser.add_argument('--model_path', type=str, default='saved_model/deep_crossing_model', help="模型文件路径")
    parser.add_argument('--embed_dim', type=int, default=8, help="嵌入维度")
    parser.add_argument('--sample_num', type=int, default=10, help="读取的样本数量")
    parser.add_argument('--batch_size', type=int, default=10, help="批量大小")
    parser.add_argument('--profile_num', type=int, default=10, help="简历数量")

    # 解析命令行参数
    args = parser.parse_args()

    # 调用预测方法
    predict_with_model(test_file=args.test_file,
                       model_path=args.model_path,
                       embed_dim=args.embed_dim,
                       sample_num=args.sample_num,
                       batch_size=args.batch_size,
                       profile_num=args.profile_num)


if __name__ == '__main__':
    main()
